import numpy as np
from torch.utils.data import DataLoader
from data import NERDataset, build_corpus
from model.config import TrainingConfig
from utils import load_model, Metrics, find_key, plot_confusion_matrix



PATH = './ckpts/bi_lstm.pkl'


def output_scores(evaluation, tag2id):
    """将结果以表格形式出来"""
    # 打印表头
    header_format = '{:>9s}  {:>9} {:>9} {:>9} {:>9}'
    header = ['precision', 'recall', 'f1-score', 'support']
    print(header_format.format('', *header))

    row_format = '{:>9s}  {:>9.4f} {:>9.4f} {:>9.4f} {:>9}'
    # 打印每个标签的 精确率、召回率、f1分数
    for tag in evaluation.tagset:
        print(row_format.format(
            find_key(tag2id, tag),
            evaluation.precision_scores[tag],
            evaluation.recall_scores[tag],
            evaluation.f1_scores[tag],
            evaluation.golden_tags_counter[tag]
        ))
    evaluation.report_avg()    # 计算并打印平均值


def output_confusion_matrix(evaluation, tag2id):
    """计算混淆矩阵"""
    print("\nConfusion Matrix:")
    tag_list = list(evaluation.tagset)
    # 初始化混淆矩阵 matrix[i][j]表示第i个tag被模型预测成第j个tag的次数
    tags_size = len(tag_list)
    matrix = []
    for i in range(tags_size):
        matrix.append([0] * tags_size)

    # 遍历tags列表
    for golden_tag, predict_tag in zip(evaluation.golden_tags, evaluation.predict_tags):
        try:
            row = tag_list.index(golden_tag)
            col = tag_list.index(predict_tag)
            matrix[row][col] += 1
        except ValueError:  # 有极少数标记没有出现在golden_tags，但出现在predict_tags，跳过这些标记
            continue

    # 输出矩阵
    row_format_ = '{:>7} ' * (tags_size + 1)
    head = [find_key(tag2id, i) for i in tag_list]
    print(row_format_.format("", *head))

    for i, row in enumerate(matrix):
        print(row_format_.format(head[i], *row))

    confu_matrix = np.array(matrix)
    plot_confusion_matrix(confu_matrix, head, "Confusion Matrix", save=False)


if __name__ == "__main__":

    # 获取测试集数据
    _, _, word_2_idx, tag_2_idx = build_corpus("train", make_vocab=True)
    test_data, test_tag = build_corpus("test", make_vocab=False)

    test_dataset = NERDataset(test_data, test_tag, word_2_idx, tag_2_idx)
    test_dataloader = DataLoader(test_dataset, TrainingConfig.dev_batch_size, shuffle=False,
                                 collate_fn=test_dataset.pro_batch_data)

    print("加载并评估Bi-LSTM模型...")
    model = load_model(PATH)
    model.lstm.flatten_parameters()  # remove warning
    model.eval()

    test_pre = []
    test_tag = []
    for test_batch_data, test_batch_tag, test_batch_len in test_dataloader:
        test_loss = model.forward(test_batch_data, test_batch_len, test_batch_tag)
        test_pre.extend(model.pre.detach().cpu().numpy().tolist())
        test_tag.extend(test_batch_tag.detach().cpu().numpy().reshape(-1).tolist())

    metrics = Metrics(test_tag, test_pre)
    output_scores(metrics, tag_2_idx)
    output_confusion_matrix(metrics, tag_2_idx)










